package com.zhang.flink.example;

import com.zhang.flink.bean.ClickEvent;
import com.zhang.flink.example.source.ClickSource;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.triggers.Trigger;
import org.apache.flink.streaming.api.windowing.triggers.TriggerResult;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;

/**
 * @title: 统计一天的pv但是要求每个小时输出一次截止当前的统计结果
 * @author: zhang
 * @date: 2022/3/22 22:47
 */
public class TriggerExample {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        env
                .addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<ClickEvent>forMonotonousTimestamps()
                                .withTimestampAssigner(new SerializableTimestampAssigner<ClickEvent>() {
                                    @Override
                                    public long extractTimestamp(ClickEvent element, long recordTimestamp) {
                                        return element.ts;
                                    }
                                })
                )
                .keyBy(r -> 1)
                .window(TumblingEventTimeWindows.of(Time.seconds(10)))
                .trigger(new MyTrigger())
                .aggregate(new CountAgg(), new WindowResult())
                .print();


        env.execute();
    }

    public static class MyTrigger extends Trigger<ClickEvent, TimeWindow> {
        //每来一条数据执行一次
        @Override
        public TriggerResult onElement(ClickEvent element, long timestamp, TimeWindow window, TriggerContext ctx) throws Exception {
            //状态作用范围当前窗口
            ValueState<Boolean> isFirst = ctx.getPartitionedState(
                    new ValueStateDescriptor<Boolean>("is-fisrt", Types.BOOLEAN)
            );
            //如果属于一个窗口的第一条数据到来，注册一个当前数据时间下一个整数秒的事件时间定时器
            if (isFirst.value() == null) {
                // 1234ms + 1000 - 1234 % 1000 = 2s
                long nextSecond = element.ts + 1000L - element.ts % 1000L;
                ctx.registerEventTimeTimer(nextSecond);
                isFirst.update(true);
            }
            //什么都不做
            return TriggerResult.CONTINUE;
        }

        @Override
        public TriggerResult onProcessingTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
            return null;
        }

        //当事件时间到达`time`是触发onEventTime执行
        @Override
        public TriggerResult onEventTime(long time, TimeWindow window, TriggerContext ctx) throws Exception {
            if (time < window.getEnd()) {
                //触发窗口计算但是不销毁窗口
                if (time + 1000L < window.getEnd()) {
                    //继续注册窗口时间范围内的每个整数秒的定时器
                    ctx.registerEventTimeTimer(time + 1000L);
                }
                return TriggerResult.FIRE;
            }
            //如果大于窗口结束时间会触发默认触发器，这里不做任何操作
            return TriggerResult.CONTINUE;
        }

        //窗口销毁的时候调用
        @Override
        public void clear(TimeWindow window, TriggerContext ctx) throws Exception {
            //状态作用范围当前窗口
            ValueState<Boolean> isFirst = ctx.getPartitionedState(
                    new ValueStateDescriptor<Boolean>("is-fisrt", Types.BOOLEAN)
            );
            isFirst.clear();
        }
    }


    public static class CountAgg implements AggregateFunction<ClickEvent, Long, Long> {
        @Override
        public Long createAccumulator() {
            return 0L;
        }

        @Override
        public Long add(ClickEvent value, Long accumulator) {
            return accumulator + 1L;
        }

        @Override
        public Long getResult(Long accumulator) {
            return accumulator;
        }

        @Override
        public Long merge(Long a, Long b) {
            return a + b;
        }
    }

    public static class WindowResult extends ProcessWindowFunction<Long, String, Integer, TimeWindow> {
        @Override
        public void process(Integer integer, ProcessWindowFunction<Long, String, Integer, TimeWindow>.Context context, Iterable<Long> elements, Collector<String> out) throws Exception {
            out.collect("窗口：" + new Timestamp(context.window().getStart()) + "~" + new Timestamp(context.window().getEnd()) + "" +
                    "中的pv是：" + elements.iterator().next() + "-----"+new Timestamp(context.currentWatermark()));
        }
    }
}
